package com.wkm.chainchain.pbft.p2p.handler;

import cn.hutool.core.collection.ConcurrentHashSet;
import com.alibaba.fastjson.JSON;
import com.wkm.chainchain.blockchain.BlockChainOp;
import com.wkm.chainchain.blockchain.block.Unit;
import com.wkm.chainchain.consensusutils.RSAUtils;
import com.wkm.chainchain.pbft.common.PBFTMsgType;
import com.wkm.chainchain.pbft.msg.PBFTMsg;
import com.wkm.chainchain.pbft.node.AllNodeInfo;
import com.wkm.chainchain.pbft.node.MsgCollection;
import com.wkm.chainchain.pbft.node.Node;
import com.wkm.chainchain.pbft.p2p.PBFTClient;
import lombok.extern.slf4j.Slf4j;

import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.UUID;

/**
 * PBFT算法的主要逻辑 在server端收到消息后根据消息状态码进行处理 并在处理成功后改变消息状态广播出去
 * TODO 目前还存在运行时某个节点断开 无法将所有共识完成 可能是因为断开的节点还有线程没有结束 导致一些消息已经记录在set中
 * reply消息却没来得及发出去
 */
@Slf4j
public class BlockchainMsgHandler {
    private static final BlockchainMsgHandler msgHandler = new BlockchainMsgHandler();

    public static BlockchainMsgHandler getInstance() {
        return msgHandler;
    }

    public void handler(String msg) {
        PBFTMsg pbftMsg = JSON.parseObject(msg, PBFTMsg.class);
        Node node = Node.getInstance();
        switch (Objects.requireNonNull(PBFTMsgType.getByCode(pbftMsg.getMsgType()))) {
            //如果是初始化消息, 将该节点放入当前服务端节点列表
            case INIT:
                init(pbftMsg);
                break;
            case REQUEST:
                //request请求只有主节点处理
                if (node.getIndex() == 0) {
                    request(pbftMsg);
                }
                break;
            case COMMIT:
                commit(pbftMsg);
                break;
            case PREPARE:
                prepare(pbftMsg);
                break;
            case PRE_PREPARE:
                //pre-prepare消息主节点不需要处理, 主节点也不会收到这类消息, 因为只有主节点自己发
                prePrepare(pbftMsg);
                break;
            case CLIENT_REPLAY:
                clientReply(pbftMsg);
                break;
            default:
        }
    }

    public void init(PBFTMsg pbftMsg) {
        Node node = pbftMsg.getNode();
        AllNodeInfo.P2P_NODE_ARRAYS[node.getIndex()] = node;
        //System.out.println(Arrays.toString(AllNodeInfo.P2P_NODE_ARRAYS));
    }

    /**
     * 主节点收到request消息 改变消息信息将其广播置P2P网络中
     */
    private void request(PBFTMsg pbftMsg) {
        //TODO 主节点收到客户端发来的消息也是需要验证的, 没有写这块逻辑 默认收到的消息是可信的
        //主节点收到pre-prepare的消息的同时, 也要将自己的prepare消息广播出去
        //在当前主节点的votePrepareMap中加入自己的共识结果
        Map<Long, Set<Integer>> votePrepareMap = MsgCollection.votePrepareMap;
        Set<Integer> msgSet = votePrepareMap.getOrDefault(pbftMsg.getMsgIndex(), new ConcurrentHashSet<>());
        msgSet.add(Node.getInstance().getIndex());
        votePrepareMap.putIfAbsent(pbftMsg.getMsgIndex(), msgSet);

        //同时将PRE_PREPARE和PREPARE消息发送出去
        setPBFTMsgInfoToNextPhase(pbftMsg, PBFTMsgType.PRE_PREPARE);
        PBFTClient.getInstance().broadcast(JSON.toJSONString(pbftMsg));
        pbftMsg.setMsgType(PBFTMsgType.PREPARE.getCode());
        PBFTClient.getInstance().broadcast(JSON.toJSONString(pbftMsg));
    }

    /**
     * 收到客户端发来的prePrepare消息, 将其保存在SET中, 通过验证则更改消息状态向全网节点广播
     */
    public void prePrepare(PBFTMsg pbftMsg) {
        log.info("收到[{}]发来的[{}]号pre-prepare消息", pbftMsg.getNodeIndex(), pbftMsg.getMsgIndex());
        //TODO (消息摘要被改变 || 请求编号不在高低水位之间 || 该消息与之前收到的相同编号消息不一致 || 已在成功池中) ====>不接受
        if (reject(pbftMsg)) {
            log.info("消息:[{}] 校验不通过 丢弃消息", pbftMsg.getMsgIndex());
            return;
        }

        /* //暂时不需要这个集合
        MsgCollection.votePrePrepareSet.add(pbftMsg.getMsgIndex());*/

        // <PREPARE, V, N, D(m), i>改变当前消息的 状态/私钥签名/节点的编号
        setPBFTMsgInfoToNextPhase(pbftMsg, PBFTMsgType.PREPARE);
        PBFTClient.getInstance().broadcast(JSON.toJSONString(pbftMsg));

        /*
        在当前节点的votePrepareMap中加入自己的共识结果
        加入map之后也需要判断一下是不是达到条件了 如果达到了 需要发送下一阶段的消息
        因为可能先收到别的节点的下一轮消息 而上一轮消息现在才收到
        */
        prepareSolve(pbftMsg, Node.getInstance().getIndex());
    }

    public void prepare(PBFTMsg pbftMsg) {
        log.info("收到[{}]节点发来的[{}]号prepare消息", pbftMsg.getNodeIndex(), pbftMsg.getMsgIndex());
        //TODO (消息摘要被改变 || 请求编号不在高低水位之间 || 该消息与之前收到的相同编号消息不一致 || 已在成功池中) ====>不接受

        if (reject(pbftMsg, MsgCollection.preparedSet)) {
            log.info("消息: [{}]<prepare>阶段校验不通过 丢弃消息", pbftMsg.getMsgIndex());
            return;
        }

        prepareSolve(pbftMsg, pbftMsg.getNodeIndex());

        /*
        在commitMap中加入自己的验证通过记录
        注意这个逻辑要写在prepareSolve后, 因为经过上面的处理, pbftMsg中node节点的Index已经从发送过来的节点序号改为为当前节点的序号了
        所以此时commitSolve的参数也可以设置为pbftMsg.getNodeIndex(), 这里为了可读性仍写成获取当前节点的Node.getInstance().getIndex()
        如果写反了的话 上面会获取不到发送节点的Index
        */
        commitSolve(pbftMsg, Node.getInstance().getIndex());
    }

    public void prepareSolve(PBFTMsg pbftMsg, int nodeIndex) {
        Long msgIndex = pbftMsg.getMsgIndex();

        Map<Long, Set<Integer>> votePrepareMap = MsgCollection.votePrepareMap;
        Set<Integer> set = votePrepareMap.getOrDefault(msgIndex, new ConcurrentHashSet<>());
        //将发送消息的node节点序号加入set
        set.add(nodeIndex);
        votePrepareMap.put(msgIndex, set);

        //判断该条消息是否达到prepare最小共识数2f + 1, 若达到则进入commit阶段
        if (set.size() >= AllNodeInfo.getLeastAgreeNodesAmount()) {
            //成功加入成功池
            MsgCollection.preparedSet.add(msgIndex);
            //在map中移除已经成功的消息释放内存
            votePrepareMap.remove(msgIndex);
            //改变消息信息
            setPBFTMsgInfoToNextPhase(pbftMsg, PBFTMsgType.COMMIT);
            //向所有节点广播commit消息
            PBFTClient.getInstance().broadcast(JSON.toJSONString(pbftMsg));
        }
    }

    public void commit(PBFTMsg pbftMsg) {
        log.info("收到[{}]节点发来的[{}]号commit消息", pbftMsg.getNodeIndex(), pbftMsg.getMsgIndex());
        //TODO (消息摘要被改变 || 请求编号不在高低水位之间 || 该消息与之前收到的相同编号消息不一致 || 已在成功池中) ====>不接受

        if (reject(pbftMsg, MsgCollection.CommittedSet)) {
            log.info("消息: [{}]<commit>阶段校验不通过 丢弃消息", pbftMsg.getMsgIndex());
            return;
        }
        commitSolve(pbftMsg, pbftMsg.getNodeIndex());
    }


    public void commitSolve(PBFTMsg pbftMsg, int nodeIndex) {
        Map<Long, Set<Integer>> voteCommitMap = MsgCollection.voteCommitMap;
        Long msgIndex = pbftMsg.getMsgIndex();

        Set<Integer> set = voteCommitMap.getOrDefault(msgIndex, new ConcurrentHashSet<>());
        set.add(nodeIndex);
        voteCommitMap.put(msgIndex, set);

        //判断该条消息是否达到最小共识数 2f + 1, 若达到则进入reply阶段
        if (set.size() >= AllNodeInfo.getLeastAgreeNodesAmount()) {
            //成功加入成功池
            MsgCollection.CommittedSet.add(msgIndex);
            //在map中移除已经成功的消息释放内存
            voteCommitMap.remove(msgIndex);
            log.info("消息[{}]达成共识", pbftMsg.getMsgIndex());
            setPBFTMsgInfoToNextPhase(pbftMsg, PBFTMsgType.CLIENT_REPLAY);
            //向客户端发送共识成功消息
            PBFTClient.getInstance().getRequestClientSockets().send(JSON.toJSONString(pbftMsg));
        }
    }


    public void clientReply(PBFTMsg pbftMsg) {
//        log.info("收到[{}]节点发来的[{}]号reply消息", pbftMsg.getNodeIndex(), pbftMsg.getMsgIndex());

        //TODO (消息摘要被改变 || 请求编号不在高低水位之间 || 该消息与之前收到的相同编号消息不一致 || 已在成功池中) ====>不接受
        if (reject(pbftMsg, MsgCollection.repliedSet)) {
//            log.info("消息: [{}]<reply>阶段校验不通过 丢弃消息", pbftMsg.getMsgIndex());
            return;
        }
        Map<Long, Set<Integer>> voteReplyMap = MsgCollection.voteReplyMap;
        Long msgIndex = pbftMsg.getMsgIndex();
        int nodeIndex = pbftMsg.getNodeIndex();

        //可能有几个线程处于当前同一位置 用ConcurrentHashMap应该没问题
        Set<Integer> set = voteReplyMap.getOrDefault(msgIndex, new ConcurrentHashSet<>());
        set.add(nodeIndex);
        voteReplyMap.put(msgIndex, set);

        //判断该条消息是否达到最小共识数f, 若达到则记录下来, 可能有两个线程不同node传来的同一消息同时进入这个if
        if (set.size() >= AllNodeInfo.getMaxFaultNodesAmount() + 1) {
            //加锁防止出现线程安全问题 这边已经加了synchronized锁, 不需要将list和set设为线程安全的集合了
            synchronized (MsgCollection.repliedSet) {
                if (MsgCollection.repliedSet.contains(msgIndex)) {
//                    log.error("消息: [{}] 出现线程问题, 已丢弃", msgIndex);
                    return;
                }
                //成功加入成功池
                MsgCollection.repliedSet.add(msgIndex);
                //在map中移除已经成功的消息释放内存
                voteReplyMap.remove(msgIndex);
//                log.info("消息[{}]完成共识已返回client, 打包生成区块", msgIndex);
                //记录已经共识成功的消息
                MsgCollection.reachedConsensusMsgList.add(msgIndex);

                Unit unit = new Unit();
                unit.setFileAddr(UUID.randomUUID().toString());
                unit.setFileHash(pbftMsg.getMessage());
                BlockChainOp.insertData(unit);
            }
        }
    }

    /**
     * 对收到的消息进行验证 判断是否需要拒绝当前消息
     * 需要验证数据是否已经在池中
     *
     * @param pbftMsg     消息
     * @param successPool 成功池
     */
    private boolean reject(PBFTMsg pbftMsg, Set<Long> successPool) {
        //签名验证成功 && 不在当前状态的成功池中
        return !(RSAUtils.verify(pbftMsg.getDigest(), pbftMsg.getMessage(),
                AllNodeInfo.P2P_NODE_ARRAYS[pbftMsg.getNodeIndex()].getPublicKey()) && !successPool.contains(pbftMsg.getMsgIndex()));
    }

    /**
     * 对收到的消息进行验证 判断是否需要拒绝当前消息
     * 只需要验证签名
     *
     * @param pbftMsg 消息
     */
    private boolean reject(PBFTMsg pbftMsg) {
        return !(RSAUtils.verify(pbftMsg.getDigest(), pbftMsg.getMessage(),
                AllNodeInfo.P2P_NODE_ARRAYS[pbftMsg.getNodeIndex()].getPublicKey()));
    }

    /**
     * 改变通过验证的其他节点的消息的信息, 改变状态码/消息发送节点编号/消息摘要digest
     *
     * @param pbftMsg 消息
     * @param type    下个阶段的类型
     */
    private void setPBFTMsgInfoToNextPhase(PBFTMsg pbftMsg, PBFTMsgType type) {
        pbftMsg.setMsgType(type.getCode());
        pbftMsg.setNodeIndex(Node.getInstance().getIndex());
        pbftMsg.setDigest(RSAUtils.sign(pbftMsg.getMessage(), Node.getInstance().getPrivateKey()));
    }
}
